/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "qfunc.h"
#include "prating.c"

float
qfunc (const float*                     loga,
       unsigned int                     n_workers,
       unsigned int                     rank,
       const float*                     gamma,
       const float*                     logbeta,
       unsigned int                     n_images,
       const float*                     c,
       unsigned int                     n_labels,
       const float*                     d,
       const float*                     pz,
       const Rating*                    ratings,
       unsigned int                     n_ratings,
       const ParameterizedDistribution* priorloga,
       const Distribution*              priorgamma,
       const Distribution*              priorlogbeta,
       const Distribution*              priorc,
       const Distribution*              priord,
       float                            hq)
{
  float q = hq;

  for (unsigned int i = 0; i < n_ratings; ++i)
    {
      for (unsigned int j = 0; j < n_labels; ++j)
        {
#define logafunc(n) loga[ratings[i].worker * rank + n]
#define cfunc(n) c[n_labels * n + j]
#define dfunc(n, k) ((k == j) ? 0 : d[n_labels * n + k])

          q += pz[j + ratings[i].image * n_labels] * 
               logprating (ratings[i].rating,
                           logafunc,
                           rank,
                           logbeta[ratings[i].image],
                           cfunc,
                           dfunc,
                           n_labels);
#undef dfunc
#undef cfunc
#undef logafunc
        }
    }

  for (unsigned int j = 0; j < rank; ++j)
    {
      q += distribution_logpdf (priorgamma, gamma[j]);
    }

  for (unsigned int i = 0; i < n_workers; ++i)
    {
      for (unsigned int j = 0; j < rank; ++j)
        {
          q += parameterized_distribution_logpdf (priorloga, 
                                                  gamma[j], 
                                                  loga[i * rank + j]);
        }
    }

  for (unsigned int i = 0; i < rank; ++i)
    {
      for (unsigned int j = 0; j < n_labels; ++j)
        {
          q += distribution_logpdf (priorc, c[n_labels * i + j]);
          q += distribution_logpdf (priord, d[n_labels * i + j]);
        }
    }

  for (unsigned int i = 0; i < n_images; ++i)
    {
      q += distribution_logpdf (priorlogbeta, logbeta[i]);
    }

  return q;
}

static void
dqfunc_h (unsigned int                     total_ratings,
          float*                           dloga,
          float*                           dgamma,
          float*                           dlogbeta,
          float*                           dc,
          float*                           dd,
          const float*                     loga,
          unsigned int                     rank,
          const float*                     gamma,
          const float*                     logbeta,
          const float*                     c,
          unsigned int                     n_labels,
          const float*                     d,
          const float*                     pz,
          const Rating*                    ratings,
          unsigned int                     n_ratings,
          const Distribution*              priorgamma,
          const Distribution*              priorc,
          const Distribution*              priord)
{
  const float epsilon = 0.001; // TODO: this is cheesy
                               // would rather ensure total mass
                               // then incremental mass
                               // but that requires modifying the e-step

  float rfrac = 
    (n_ratings == total_ratings) 
        ? 1.0 
        : ((float) n_ratings) / ((float) total_ratings);

  for (unsigned int i = 0; i < n_ratings; ++i)
    {
      for (unsigned int j = 0; j < n_labels; ++j)
        {
          float tmpdloga[rank];
          float tmpdc[rank];
          float tmpdd[rank * n_labels];

#define logafunc(n) loga[ratings[i].worker * rank + n]
#define cfunc(n) c[n_labels * n + j]
#define dfunc(n, k) ((k == j) ? 0 : d[n_labels * n + k])

          if (pz[j + ratings[i].image * n_labels] > epsilon)
            {
              dlogbeta[ratings[i].image] +=
                pz[j + ratings[i].image * n_labels] *
                dlogpratingdall (tmpdloga,
                                 tmpdc,
                                 tmpdd,
                                 ratings[i].rating,
                                 logafunc,
                                 rank,
                                 logbeta[ratings[i].image],
                                 cfunc,
                                 dfunc,
                                 n_labels);
#undef dfunc
#undef cfunc
#undef logafunc

              for (unsigned int n = 0; n < rank; ++n)
                {
                  dloga[ratings[i].worker * rank + n] += 
                    pz[j + ratings[i].image * n_labels] * tmpdloga[n];
                  dc[n_labels * n + j] += 
                    pz[j + ratings[i].image * n_labels] * tmpdc[n];
   
                  for (unsigned int k = 0; k < n_labels; ++k)
                    {
                      if (k != j)
                        {
                          dd[n_labels * n + k] += 
                            pz[j + ratings[i].image * n_labels] * 
                            tmpdd[n_labels * n + k];
                        }
                    }
                }
            }
        }
    }

  for (unsigned int j = 0; j < rank; ++j)
    {
      dgamma[j] += rfrac * distribution_dlogpdf (priorgamma, gamma[j]);
    }

  for (unsigned int i = 0; i < rank; ++i)
    {
      for (unsigned int j = 0; j < n_labels; ++j)
        {
          dc[n_labels * i + j] += 
            rfrac * distribution_dlogpdf (priorc, c[n_labels * i + j]);
          dd[n_labels * i + j] += 
            rfrac * distribution_dlogpdf (priord, d[n_labels * i + j]);
        }
    }
}

void
dqfunc (float*                           dloga,
        float*                           dgamma,
        float*                           dlogbeta,
        float*                           dc,
        float*                           dd,
        const float*                     loga,
        unsigned int                     n_workers,
        unsigned int                     rank,
        const float*                     gamma,
        const float*                     logbeta,
        unsigned int                     n_images,
        const float*                     c,
        unsigned int                     n_labels,
        const float*                     d,
        const float*                     pz,
        const Rating*                    ratings,
        unsigned int                     n_ratings,
        const ParameterizedDistribution* priorloga,
        const Distribution*              priorgamma,
        const Distribution*              priorlogbeta,
        const Distribution*              priorc,
        const Distribution*              priord)
{
  memset (dloga, 0, n_workers * rank * sizeof (dloga[0]));
  memset (dgamma, 0, rank * sizeof (dgamma[0]));
  memset (dlogbeta, 0, n_images * sizeof (dlogbeta[0]));
  memset (dc, 0, rank * n_labels * sizeof (dc[0]));
  memset (dd, 0, rank * n_labels * sizeof (dd[0]));

  dqfunc_h (n_ratings, dloga, dgamma, dlogbeta, dc, dd, loga, rank, 
            gamma, logbeta, c, n_labels, d, pz,
            ratings, n_ratings, priorgamma, priorc, priord);

  for (unsigned int i = 0; i < n_images; ++i)
    {
      dlogbeta[i] += distribution_dlogpdf (priorlogbeta, logbeta[i]);
    }

  for (unsigned int i = 0; i < n_workers; ++i)
    {
      for (unsigned int j = 0; j < rank; ++j)
        {
          dgamma[j] += 
            parameterized_distribution_d0logpdf (priorloga,
                                                 gamma[j],
                                                 loga[i * rank + j]);

          dloga[i * rank + j] += 
            parameterized_distribution_d1logpdf (priorloga,
                                                 gamma[j],
                                                 loga[i * rank + j]);
        }
    }
}

void
dqfunc_single (unsigned int                     the_rating,
               unsigned int                     n_image_ratings,
               unsigned int                     n_worker_ratings,
               float*                           dloga,
               float*                           dgamma,
               float*                           dlogbeta,
               float*                           dc,
               float*                           dd,
               const float*                     loga,
               unsigned int                     rank,
               const float*                     gamma,
               const float*                     logbeta,
               const float*                     c,
               unsigned int                     n_labels,
               const float*                     d,
               const float*                     pz,
               const Rating*                    ratings,
               unsigned int                     n_ratings,
               const ParameterizedDistribution* priorloga,
               const Distribution*              priorgamma,
               const Distribution*              priorlogbeta,
               const Distribution*              priorc,
               const Distribution*              priord)
{
  dqfunc_h (n_ratings, dloga, dgamma, dlogbeta, dc, dd, loga, rank, 
            gamma, logbeta, c, n_labels, d, pz, 
            ratings + the_rating, 1, priorgamma, priorc, priord);

  unsigned int i = ratings[the_rating].image;

  dlogbeta[i] += 
    (1.0 / n_image_ratings) * distribution_dlogpdf (priorlogbeta, logbeta[i]);

  unsigned int w = ratings[the_rating].worker;
  float rfrac = 1.0 / n_worker_ratings;

  for (unsigned int j = 0; j < rank; ++j)
    {
      dgamma[j] += 
        rfrac *
        parameterized_distribution_d0logpdf (priorloga,
                                             gamma[j],
                                             loga[w * rank + j]);

      dloga[w * rank + j] += 
        rfrac *
        parameterized_distribution_d1logpdf (priorloga,
                                             gamma[j],
                                             loga[w * rank + j]);
    }
}
